"""
A module to encapsulate a Likelihood source model xml file.

@author J. Chiang <jchiang@slac.stanford.edu>
@modified by Z.Q. Shen
"""
from __future__ import print_function
import os
#import string
from xml.dom import minidom
from math import pi, cos, sin, acos


bools = {'True' : 1, 'False' : 0, 'T' : 1, 'F' : 0,
         't' : 1, 'f' : 0, 'TRUE' : 1, 'FALSE' : 0,
         'true' : 1, 'false' : 0, '1' : 1, '0' : 0,
         0: 0, 1: 1, True: 1, False: 1}

d2r = pi / 180.

def mybool(Input):
    return bools[Input]


def angsep(ra1, dec1, ra2, dec2):
    ra1 *= d2r
    dec1 *= d2r
    ra2 *= d2r
    dec2 *= d2r
    diffCosine = cos(dec1) * cos(dec2) * cos(ra1-ra2) + sin(dec1) * sin(dec2)
    dC = '%.10f' % diffCosine
    return acos(float(dC)) / d2r #returns values between 0 and pi radians


class SourceModel(object):
    _BLANK_MODEL = '\n'.join(('<?xml version="1.0"?>',
                            '<source_library title="source library">',
                            '</source_library>'))

    def __init__(self, xmlFile=None, doc=None, raiseErr=True):
        if not doc is None:
            self.doc = doc
        elif xmlFile:
            if os.path.isfile(xmlFile):
                self.filename = xmlFile
                self.doc = minidom.parse(self.filename)
            elif raiseErr:
                raise IOError('file %s does not exist' % xmlFile)
            else:
                print('file %s does not exist, create a blank ones' % xmlFile)
                self.filename = xmlFile
                self.doc = minidom.parseString(self._BLANK_MODEL)
        else:
            self.filename = xmlFile
            self.doc = minidom.parseString(self._BLANK_MODEL)

        srcs = self.doc.getElementsByTagName("source")
        self.srcList, self.freeSrcs = {}, []
        try:
            for src in srcs:
                srcname = str(src.getAttribute("name"))
                self.srcList[srcname] = Source(src)
                if self.srcList[srcname].isfree: self.freeSrcs.append(srcname)
        except ValueError:
            pass

    def setAttributes(self):
        'save all the changes in attribute to minidom.dom obj'
        for src in self.srcList.values():
            src.setAttributes()

    def __getitem__(self, name):
        return self.srcList[name]

    def __setitem__(self, name, value):
        value2 = value.copy()
        if name is None:
            name = value2.name
        else:
            value2.name = name
            value2.setAttributes()
        source_library = self.doc.getElementsByTagName('source_library')[0]
        if name in self.srcList.keys():
            source_library.replaceChild(value2.node, self.srcList[name].node)
        else:
            source_library.appendChild(value2.node)
        self.srcList[name] = value2

    def __delitem__(self, name):
        source_library = self.doc.getElementsByTagName('source_library')[0]
        source_library.removeChild(self.srcList[name].node)
        del self.srcList[name]

    def addsrc(self, value):
        name = value.name
        if name in self.srcList:
            raise RuntimeError('the source %s is already exist'%name)

        value2 = value.copy()
        source_library = self.doc.getElementsByTagName('source_library')[0]
        source_library.appendChild(value2.node)
        self.srcList[name] = value2

    @property
    def names(self):
        return sorted(self.srcList.keys())

    def __repr__(self):
        return repr(self.names)

    @staticmethod
    def _indent_lines(lines):
        level = 0
        for i in range(1, len(lines)):
            if lines[i].find('<!--') != -1 and lines[i].find('-->') != -1:  # skip comment lines
                lines[i] = '    '*level + lines[i]
            elif lines[i].find('</') == -1:
                lines[i] = '    '*level + lines[i]
                if lines[i].find('/>') == -1:
                    level += 1
            else:
                level -= 1
                lines[i] = '    '*level + lines[i]

    @staticmethod
    def _cleanXml(document):
        try:
            doc = minidom.parse(document)
        except AttributeError:
            doc = document
        lines = doc.toxml().split('>')
        clean_lines = []
        for line in lines:
            if line.find('<') != -1:
                line = line + '>'
                clean_lines.append(line.strip())
        SourceModel._indent_lines(clean_lines)
        xmlString = '\n'.join(clean_lines)
        new_doc = minidom.parseString(xmlString)
        return new_doc

    def toXml(self):
        self.setAttributes()
        try:
            doc = self._cleanXml(self.doc)
        except Exception as err:
            print('[WARN] Some error occurs when making clean xml:')
            print(str(err).strip())
            doc = self.doc
        return str(doc.toxml())

    def printXml(self):
        print(self.toXml())

    def writeTo(self, filename=None):
        if filename == None:  # Overwrite the existing file.
            filename = self.filename
        self.setAttributes()
        try:
            doc = self._cleanXml(self.doc)
        except Exception as err:
            print('[WARN] Some error occurs when making clean xml:')
            print(str(err).strip())
            doc = self.doc
        fh = open(filename, 'w')
        fh.write(doc.toxml() + '\n')
        fh.close()
        self.filename = filename

    def sort(self):
        source_library = self.doc.getElementsByTagName('source_library')[0]
        for child in source_library.childNodes:
            if child.nodeType == 8:
                source_library.removeChild(child)
        srcNames = sorted(self.srcList.keys())
        for name in srcNames:
            source_library.removeChild(self.srcList[name].node)
            source_library.appendChild(self.srcList[name].node)

    def copy(self):
        self.setAttributes()
        return SourceModel(xmlFile=self.filename, doc=self.doc.cloneNode(True))

    def makeRegion(self, regfile, color1='green', color2='magenta', color3='cyan', ptype='cross'):
        """
        color = white, black, red, green, blue, cyan, magenta, yellow
        color1 for free source (nfree>1), color2 for fixed source, color3 for free source (nfree=1)
        """
        myreg = MyRegion()
        for s in self.srcList.values():
            name = s.name
            if s.spatialModel.type == 'SkyDirFunction' or \
               s.spatialModel.type == 'RadialGaussian' or \
               s.spatialModel.type == 'RadialDisk':
                if 'RA' in s.spatialModel.parameters.keys():
                    coord = 'CEL'
                    x = s.spatialModel.parameters['RA'].value
                    y = s.spatialModel.parameters['DEC'].value
                elif 'LON' in s.spatialModel.parameters.keys():
                    coord = 'GAL'
                    x = s.spatialModel.parameters['LON'].value
                    y = s.spatialModel.parameters['LAT'].value
                else:
                    raise RuntimeError('No position available!')
            else:
                try:
                    x = float(s['RA'])
                    y = float(s['DEC'])
                    coord = 'CEL'
                except:
                    print('{} has no position, skipped'.format(name))
                    continue

            nfree = 0
            for p in s.spectrum.parameters.values():
                if p.free == 1:
                    nfree += 1
            if nfree <= 0:
                color = color2
            elif nfree <= 1:
                color = color3
            else:
                color = color1
            myreg.addPoint(name, x, y, color, coord, ptype)
        myreg.writeTo(regfile)

    def freeze(self, name, pname=None):
        """
        If pname is None, freeze all the parameters of the input name.
        """
        try:
            src = self.srcList[name]
        except KeyError:
            raise IOError('Unknown srcname %s' % name)
        else:
            src.freeze(pname)

    def thaw(self, name, pname):
        """
        thaw the parameter pname of a src.
        """
        try:
            src = self.srcList[name]
        except KeyError:
            raise IOError('Unknown srcname %s' % name)
        else:
            src.thaw(pname)


class DomElement(object):

    def __init__(self, node, converter=None):
        self.node = node
        attributes = {}
        for key in node.attributes.keys():
            if converter and key in {'min', 'max', 'scale', 'value', 'error'}:
                attributes[key] = converter(node.getAttribute(key))
            elif key == "free":
                attributes[key] = mybool(node.getAttribute(key))
            else:
                attributes[key] = str(node.getAttribute(key))
        self.__dict__.update(attributes)

    def deleteChildElements(self, tagName):
        children = self.node.getElementsByTagName(tagName)
        for child in children:
            self.node.removeChild(child)

    def setAttributes(self):
        for key in self.node.attributes.keys():
            self.node.setAttribute(key, "%s" % self.__dict__[key])

    def toXml(self):
        return str(self.node.toxml())

    def printXml(self):
        print(self.toXml())

    def __getitem__(self, name):
        return self.__dict__[name]

    def __delitem__(self, name):
        del self.__dict__[name]
        self.node.removeAttribute(name)

    def __setitem__(self, key, value):
        self.__dict__[key] = value
        self.node.setAttribute(key, "%s" % self.__dict__[key])


class Source(DomElement):

    def __init__(self, node):
        DomElement.__init__(self, node)
        self.functions = {}
        (spectrum, ) = node.getElementsByTagName("spectrum")
        self.functions["spectrum"] = Function(spectrum)
        (spatialModel, ) = node.getElementsByTagName("spatialModel")
        self.functions["spatialModel"] = Function(spatialModel)
        self.__dict__.update(self.functions)
        if self.spectrum.freePars:
            self.isfree = True
        else:
            self.isfree = False

    def setAttributes(self):
        DomElement.setAttributes(self)
        for func in self.functions.values():
            func.setAttributes()

    def summary(self):
        if self.type == "PointSource":
            (line, ) = ("%s: (RA, Dec) = (%.2f, %.2f)"
                        % (self.name,
                           self.spatialModel.RA.value,
                           self.spatialModel.DEC.value), )
        else:
            line = "%s: %s" % (self.name, self.spatialModel.file.split('/')[-1])
        return line

    def angsep(self, ra, dec):
        if self.type == "PointSource":
            return angsep(self.spatialModel.RA.value, self.spatialModel.DEC.value, ra, dec)
        else:
            return None

    def copy(self):
        self.setAttributes()
        return Source(self.node.cloneNode(True))

    def freeze(self, pname=None):
        if pname is None:
            for par in self.spectrum.parameters.values():
                if hasattr(par, 'error'): del par['error']
                par.free = 0
            self.spectrum.freePars = []
        else:
            if pname in self.spectrum.parameters.keys():
                mypar = self.spectrum.parameters[pname]
                if hasattr(mypar, 'error'): del mypar['error']
                mypar.free = 0
                try:
                    self.spectrum.freePars.remove(pname)
                except ValueError:
                    pass
            else:
                raise IOError('Parameter %s is not found in src %s' % (pname, self.name))

        if not self.spectrum.freePars:
            self.isfree = False

    def thaw(self, pname):
        if pname in self.spectrum.parameters.keys():
            mypar = self.spectrum.parameters[pname]
            if hasattr(mypar, 'error'): del mypar['error']
            mypar.free = 1
            self.spectrum.freePars.append(pname)
        else:
            raise IOError('Parameter %s is not found in src %s' % (pname, self.name))

        if self.spectrum.freePars:
            self.isfree = True


class Function(DomElement):

    def __init__(self, node):
        DomElement.__init__(self, node)
        self.parameters = {}
        paramElements = node.getElementsByTagName("parameter")
        for param in paramElements:
            name = str(param.getAttribute("name"))
            self.parameters[name] = Parameter(param)
        self.__dict__.update(self.parameters)
        self.freePars = []
        for k, v in self.parameters.items():
            if v.free == 1: self.freePars.append(k)

    def setAttributes(self):
        DomElement.setAttributes(self)
        for param in self.parameters.values():
            param.setAttributes()

class Parameter(DomElement):

    def __init__(self, node):
        DomElement.__init__(self, node, float)


# Something about the region file
class MyRegion(object):
    def __init__(self, infile=None):
        if infile is None:
            self.regionData = ['# Region File format: DS9 version ?',
                               '# Created by XmlTools.SMKernel.SourceModel',
                               'global font="roman 10 normal" move =0']
        else:
            fh = open(infile, 'r')
          # self.regionData = map(string.strip, fh.readlines())
            self.regionData = [line.strip() for line in fh]
            fh.close()

    def writeTo(self, outfile):
        fh = open(outfile, 'w')
        fh.write('\n'.join(self.regionData))
        fh.close()

    @staticmethod
    def parseCoord(coord):
        _coord = coord.lower()
        if _coord == 'gal' or _coord == 'g':
            line = 'GALACTIC;'
        elif _coord == 'cel' or _coord == 'c':
            line = 'FK5;'
        else:
            raise IOError('Unknown coordsys %s' % coord)
        return line

    def addPoint(self, name, x, y, color='green', coord='CEL', ptype='cross'):
        """
        x, y in deg, coord=CEL|GAL
        """
        line = self.parseCoord(coord)
        line += 'point(%.3f,%.3f) # point = %s color = %s text={%s}' % (x, y, ptype, color, name)
        self.regionData.append(line)

    def addCircle(self, name, x, y, radius, color='green', coord='CEL'):
        line = self.parseCoord(coord)
        line += 'circle(%.3f,%.3f,%.3f) # color = %s text={%s}' % (x, y, radius, color, name)
        self.regionData.append(line)

    def addPointXSrc(self, xsrc, color='green', ptype='cross'):
        sname = xsrc['name']
        try:
            ra = xsrc.spatialModel.RA.value
            dec= xsrc.spatialModel.DEC.value
        except AttributeError:
            try:
                ra = float(xsrc['RA'])
                dec= float(xsrc['DEC'])
            except KeyError:
                raise RuntimeError('no coordinates found in the xsrc')
        self.addPoint(sname, ra, dec, color=color, coord='CEL', ptype=ptype)

    def addCircleXSrc(self, xsrc, radius, color='green'):
        sname = xsrc['name']
        try:
            ra = xsrc.spatialModel.RA.value
            dec= xsrc.spatialModel.DEC.value
        except AttributeError:
            try:
                ra = float(xsrc['RA'])
                dec= float(xsrc['DEC'])
            except KeyError:
                raise RuntimeError('no coordinates found in the xsrc')
        self.addCircle(sname, ra, dec, radius=radius, color=color, coord='CEL')
